import json
import os
import shutil
import time

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from data.neu_seg_competition_nation.dataset import Neu_Seg_Competition_Dataset_nation
from model.model_1 import *
from utils import convert_pred_to_mask, get_scores


def submit_for_ans(self_model_path):
    batch_size = 1  # 评估模式，不能动这个
    pin_memory = False
    num_workers = 1
    num_classes = 4
    # 创建设备对象
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    self_model = self_net()
    self_model.load_state_dict(torch.load(self_model_path))

    # 将模型移动到指定设备
    self_model = self_model.to(device,non_blocking=True)

    test_dataset = Neu_Seg_Competition_Dataset_nation(data_path='./data/neu_seg_competition_nation', data_type='test',seed=16788)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                             pin_memory=pin_memory)

    self_model.eval()
    shutil.copy2('./model/model_1.py', './国赛提交目录/model.py') #复制自建模型结构
    shutil.copy2(self_model_path, './国赛提交目录/model.pth')    #复制自建模型的参数
    self_model_t_all = []
    self_total_u = torch.zeros(num_classes, device=device)
    self_total_n = torch.zeros(num_classes, device=device)
    with torch.no_grad():
        for data in tqdm(test_loader):
            image, mask, image_names = data
            image_name = image_names[0]
            image = image.to(device)
            mask = mask.to(device)
            self_model_tstart = time.time()
            self_pred = self_model(image)
            self_model_tstop = time.time()
            self_mask = convert_pred_to_mask(self_pred)
            self_ni, self_ui = get_scores(self_pred, mask, device)
            self_ni = self_ni.to(device)
            self_ui = self_ui.to(device)
            self_total_n += self_ni
            self_total_u += self_ui
            self_mask_np = self_mask.cpu().numpy()
            mask_np = mask.cpu().numpy()
            self_model_t_all.append(self_model_tstop-self_model_tstart)

    self_total_u = torch.where(self_total_u < 1, 1, self_total_u)
    self_classiou = self_total_n / self_total_u
    self_miou = (self_classiou[1].item() + self_classiou[2].item() + self_classiou[3].item()) / 3
    self_fps = 1/(np.mean(self_model_t_all))
    self_parameters = sum(p.numel() for p in self_model.parameters() if p.requires_grad)
    composite_dict = {
        "OursModel": {
            "Class1_IoU": self_classiou[1].item(),
            "Class2_IoU": self_classiou[2].item(),
            "Class3_IoU": self_classiou[3].item(),
            "mIoU": self_miou,
            "FPS": self_fps,
            "Parameters": self_parameters
        }
    }
    with open('./国赛提交目录/关键指标数据文档.txt', 'w') as f:
        json.dump(composite_dict, f)
    formatted_json = json.dumps(composite_dict, indent=4)
    print(formatted_json)

if __name__ == '__main__':
    self_model_path = './save_model_best/seed_4204_冲榜了_miou_0.8269644180933634_model_best_epoch_231.pth'
    submit_for_ans(self_model_path)
    print('现在 国赛提交目录有了一部分文件了 ，你 现在可以去 ./国赛提交目录/npy生成文件材料/ 运行run.py 然后如果提示你可以提交赛克了，就说明一切就绪。')












